package com.yxdj.ai.config;

import dev.langchain4j.data.document.Document;
import dev.langchain4j.data.document.loader.FileSystemDocumentLoader;
import dev.langchain4j.data.embedding.Embedding;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.memory.chat.ChatMemoryProvider;
import dev.langchain4j.memory.chat.MessageWindowChatMemory;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.model.ollama.OllamaEmbeddingModel;
import dev.langchain4j.model.output.Response;
import dev.langchain4j.rag.content.retriever.EmbeddingStoreContentRetriever;
import dev.langchain4j.store.embedding.EmbeddingStore;
import dev.langchain4j.store.embedding.inmemory.InMemoryEmbeddingStore;
import org.springframework.boot.CommandLineRunner;
import org.springframework.boot.SpringBootConfiguration;
import org.springframework.context.annotation.Bean;

import java.util.List;

/**
 * @author 韩总
 */
@SpringBootConfiguration
public class LangChai4jAiConfiguration {

//    @Bean
//    public ChatMemory langchain4jChatMemory() {
//        MessageWindowChatMemory messageWindowChatMemory = MessageWindowChatMemory.withMaxMessages(100);
//        return messageWindowChatMemory;
//    }

    @Bean
    public ChatMemoryProvider chatMemoryProvider() {
       return memoryId-> MessageWindowChatMemory.builder().id(memoryId).maxMessages(100).build();
    }

    @Bean
    public EmbeddingStore<TextSegment> embeddingStore(){
        return new InMemoryEmbeddingStore<>();
    }

    //因为在IOC容器中已经存在一个名字为ollamaEmbeddingModel的bean，（Spring AI提供的）所以这里只能通过这种方式来配置嵌入模型
    @Bean
    public EmbeddingModel embeddingModel(){
        return OllamaEmbeddingModel.builder()
                .baseUrl("http://127.0.0.1:11434")
                .modelName("qwen:4b")
                .build();
    }

    //配置嵌入存储检索内容的对象
    @Bean
    public EmbeddingStoreContentRetriever embeddingStoreRetriever(EmbeddingModel embeddingModel, EmbeddingStore<TextSegment> embeddingStore){
        return EmbeddingStoreContentRetriever.builder()
                //配置嵌入模型
                .embeddingModel(embeddingModel)
                //配置向量数据库
                .embeddingStore(embeddingStore)
                //配置最大取前几个结果
                .maxResults(5)
                //配置能容忍的最低评分
                .minScore(0.3)
                .build();
    }

    @Bean
    CommandLineRunner  commandLineRunner(EmbeddingModel embeddingModel ,EmbeddingStore embeddingStore){
        return args -> {
            Document document = FileSystemDocumentLoader.loadDocument("G:\\代码\\cloud-repair-home\\yxdj-server\\yxdj-modules\\yxdj-ai\\src\\main\\resources\\rag\\scope.txt");
            //将文档转换为文本片段
            TextSegment textSegment = document.toTextSegment();
            //嵌入
            Response<Embedding> embed = embeddingModel.embed(textSegment);
            //获取嵌入向量
            Embedding content = embed.content();
            //将向量保存到向量数据库
            embeddingStore.add(content,textSegment);
            //将向量转换为集合
            List<Float> floats = content.vectorAsList();
            System.out.println("floats = " + floats);
            //获取向量维度
            int dimension = content.dimension();
            System.out.println("dimension = " + dimension);
        };
    }
}
